import ollama
import json


# 1. 定义工具函数
def add_numbers(num1: float, num2: float):
    return {
        "result": num1 + num2,
        "num1": num1,
        "num2": num2
    }

# 2. 定义工具描述(schema)
TOOLS = [
    {
        "type": "function",
        "function": {
            "name": "add_numbers",
            "description": "每当你需要计算两个数字的和时，都尽可能使用这个工具。",
            "parameters": {
                "type": "object",
                "properties": {
                    "num1": {
                        "type": "number",
                        "description": "第一个数字"
                    },
                    "num2": {
                        "type": "number",
                        "description": "第二个数字"
                    }
                },
                "required": ["num1", "num2"]
            }
        }
    }
]

# 3. 实现工具函数调用（基于函数名）
def get_function_by_name(name):
    # 只调用mapping以内的函数，以防止执行任意代码（例如：黑客会构建某些不安全的函数名）
    function_mapping = {
        "add_numbers": add_numbers
    }
    return function_mapping.get(name)

def process_llm_func_call(response):
    """处理大模型的func call响应"""
    if response.tool_calls:  
        for tool_call in response.tool_calls:
            func_name = tool_call.function.name
            func_args = tool_call.function.arguments
            
            # 执行函数调用：根据函数名执行函数调用（通过映射）
            func = get_function_by_name(func_name)
            result = func(**func_args)
            return result

# 4. 实现llm
client = ollama.Client(host='http://192.168.20.43:11434')
def llm(messages):
    response = client.chat( # 仍然是用chat来调用
        model='qwen2.5:14b', 
        tools=TOOLS,    # 添加工具列表
        messages=messages
    )
    return response.message


# 5. 实现主流程，包括messages记录，以及while True的对话循环
def main():
    """主要执行流程"""
    # 1. 准备初始对话历史 
    messages = [ 
        {"role": "system", "content": "你是一位友好的AI助手。"},
    ]
    
    print("欢迎使用AI助手，请输入您的问题(输入'quit'退出):")
    
    while True:
        # 获取用户输入
        user_input = input("> ")
        if user_input.lower() == 'quit':
            print("感谢使用，再见!")
            break
            
        # 添加用户输入到对话历史
        messages.append({"role": "user", "content": user_input})
        
        # 调用大模型,传入工具描述 
        response = llm(messages)
        # 处理响应：工具调用
        
        # 打印调试信息
        if hasattr(response, 'tool_calls') and response.tool_calls:
            result = process_llm_func_call(response)
            

            # 将工具调用结果加入对话历史 
            messages.append({ 
                "role": "assistant", 
                "content": None, 
                "tool_calls": response.tool_calls 
            }) 
            print("[TOOL]:", response.tool_calls)
            
            messages.append({ 
                "role": "tool", 
                #"tool_call_id": response.tool_calls[0].id, 
                "content": json.dumps(result) 
            }) 
            print("[TOOL]:", json.dumps(result) )

            # 继续对话 
            final_response = llm(messages=messages)
            print("[AI]:", final_response.content)
            messages.append({"role": "assistant", "content": final_response.content})
        else:
            # 直接输出响应
            print("[AI]:", response.content)
            messages.append({"role": "assistant", "content": response.content})

if __name__ == "__main__":
    main()